#include "ball_query_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
    constexpr int32_t BUFFER_NUM = 2;
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t CALC_SIZE = 256; 
    uint32_t RoundUp(uint32_t num, uint32_t rnd)
    {
        if (rnd == 0) {
            return 0;
        }
        return (num + rnd - 1) / rnd * rnd;
    }
    uint32_t RoundDown(uint32_t num, uint32_t rnd)
    {
        if (rnd == 0) {
            return 0;
        }
        return (num) / rnd * rnd;
    }     
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
  BallQueryTilingData tiling;
    auto ascendcPlatform = platform_ascendc:: PlatformAscendC(context->GetPlatformInfo());
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    const float* pMin = (context->GetAttrs()->GetFloat(0));
    const float* pMax = (context->GetAttrs()->GetFloat(1));
    const int64_t* pNum = (context->GetAttrs()->GetInt(2));
    const gert::Tensor * x1 = context->GetInputTensor(0);
    const gert::Tensor * x2 = context->GetInputTensor(1);
    const gert::Tensor * batch = context->GetInputTensor(2);
    auto x1_shape = x1->GetOriginShape();
    uint32_t x1_dimNum = x1_shape.GetDimNum(); 
    auto x2_shape = x2->GetOriginShape();
    uint32_t tilingKey = 0;
    uint32_t B, N, M;
    uint32_t batch_shape;
    if(x1_dimNum == 3) {
        tilingKey = 1;
        B = x1_shape.GetDim(0);
        M = x1_shape.GetDim(1);
        N = x2_shape.GetDim(1);
    }
    else if(x1_dimNum == 2) {
        tilingKey = 3;
        M = x1_shape.GetDim(0); 
        N = x2_shape.GetDim(0);
        batch_shape = batch->GetShapeSize();
    }
    uint32_t dTypeSize;
    auto x_dtype = x1->GetDataType();  
    if((x_dtype == ge::DT_FLOAT16))
    {
        dTypeSize = 2; 
        tilingKey += 1;
    }
    else if((x_dtype == ge::DT_FLOAT))
    {
        dTypeSize = 4; 
    }
    else
    {
    }   
    context->SetTilingKey(tilingKey);
    uint32_t MAlign = RoundUp(M, (BLOCK_SIZE/dTypeSize));
    uint32_t NAlign = RoundUp(N, (BLOCK_SIZE/dTypeSize));
    uint32_t sample_numAlign = RoundUp((*pNum), (BLOCK_SIZE/sizeof(int32_t)));
    context->SetBlockDim(aivNum);
    uint32_t max_baseB = 1;
    if(tilingKey == 1){
        max_baseB = 2;
    }
    uint32_t baseB;
    uint32_t B_corAlign = RoundUp(B, aivNum); 
    if((B_corAlign / aivNum) < max_baseB)
    {
        baseB = (B_corAlign / aivNum);
    }    
    else
    {
        baseB = max_baseB;
    }
    tiling.set_B(B);
    tiling.set_N(N);
    tiling.set_M(M);
    tiling.set_NAlign(NAlign);
    tiling.set_MAlign(MAlign);
    tiling.set_neg_min((*pMin) * (*pMin));
    tiling.set_neg_max((*pMax) * (*pMax));
    tiling.set_sample_num(*pNum);
    tiling.set_sample_numAlign(sample_numAlign);
    tiling.set_baseB(baseB);
    tiling.set_batch_shape(batch_shape);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    return ge::GRAPH_SUCCESS;
}
}
namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}
namespace ops {
class BallQuery : public OpDef {
public:
    explicit BallQuery(const char* name) : OpDef(name)
    {
        this->Input("xyz")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("center_xyz")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("xyz_batch_cnt")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("center_xyz_batch_cnt")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("idx")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("min_radius").Float();
        this->Attr("max_radius").Float();
        this->Attr("sample_num").Int();
        this->SetInferShape(ge::InferShape);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};
OP_ADD(BallQuery);
}
